#!/usr/bin/env bash

# NUMA node and CPU list files
# /sys/devices/system/node/possible
# /sys/devices/system/node/online
# /sys/devices/system/node/node[0-9]*/cpulist
# /sys/devices/system/cpu/possible
# /sys/devices/system/cpu/online

# NUMA node of net, block or pci device
# /sys/class/net/<net dev name>/device/numa_node
# /sys/class/block/<nvme dev name>/device/device/numa_node
# /sys/devices/pci[0-9a-fA-F:/]+\.[0-9]+/numa_node

NODE_ARRAY=()
CPULIST_ARRAY=()

# example: "0 2 4 1 3 2" --> "0 1 2 2 3 4"
function sort_numbers()
{
	local result=($(echo $@ | tr ' ' '\n' | sort -n | tr '\n' ' '))
	echo ${result[@]}
}

# example: "0,2-4,6..7" --> "0 2 3 4 6 7"
function expand_range()
{
	local range=${@//,/ }
	range=${range//-/..}
	local result=""
	local r
	for r in ${range}; do
		# eval lets us use vars in {#..#} range
		[[ ${r} =~ '..' ]] && r="$(eval echo {${r}})"
		result+=" ${r}"
	done
	echo ${result}
}

# sort after expand a range
function expand_range_ex()
{
	echo $(sort_numbers $(expand_range ${@}))
}

# example: "0, 2, 3 4 6 7" --> "0,2-4,6-7"
function compact_range()
{
	local range=${@//,/ }
	local numbers=(${range})
	local result=""
	local begin
	local idx

	# is empty range?
	((${#numbers[@]} == 0)) && echo ${result} && return

	# save all ranges but the last one
	begin=${numbers[0]}
	for ((idx = 1; idx < ${#numbers[@]}; idx++)); do
		local current=${numbers[idx]}
		local prev=${numbers[idx-1]}
		((${current} == $((prev + 1)))) && continue
		if ((${begin} == ${prev})); then
			result+="${begin},"
		else
			result+="${begin}-${prev},"
		fi
		begin=${current}
	done

	# save the last range
	local last=${numbers[-1]}
	if ((${begin} == ${last})); then
		result+="${begin}"
	else
		result+="${begin}-${last}"
	fi

	echo ${result}
}

# expand and sort before compact a range
function compact_range_ex()
{
	echo $(compact_range $(sort_numbers $(expand_range ${@})))
}

# input:  [count(0)] [begin(0)] [step(1)]
# output: [compact index range]
# other method: echo {0..5}, echo {0..10..2} or sed, e.g.
function index_range()
{
	local count=${1:-0}
	local begin=${2:-0}
	local step=${3:-1}
	local idx=${begin}
	local cnt
	local result
	for ((cnt = 0; ${cnt} < ${count}; cnt++)); do
		result=${result:+${result},}${idx}
		((idx += ${step}))
	done
	echo ${result}
}

# input:  [begin(0)] [max offset(0)] [step(1)]
# output: <count> [compact index range]
# other method: echo {0..5}, echo {0..10..2} or sed, e.g.
function index_range_ex()
{
	local begin=${1:-0}
	local max_offset=${2:-0}
	local step=${3:-1}
	local end=$((${begin} + ${max_offset}))
	local cnt=0
	local idx
	local result
	((${step} <= 0)) && echo "0" && return
	for ((idx = ${begin}; idx <= ${end}; idx += ${step})); do
		result=${result:+${result},}${idx}
		((cnt++))
	done
	echo "${cnt} ${result}"
}

# read NUMA node and CPU list
function init_cpu_nodes()
{
	NODE_ARRAY=()
	CPULIST_ARRAY=()

	NODE_ARRAY=($(expand_range `cat /sys/devices/system/node/online`))
	local node
	for node in "${NODE_ARRAY[@]}"; do
		local cpulist=$(expand_range `cat /sys/devices/system/node/node${node}/cpulist`)
		CPULIST_ARRAY[${node}]=${cpulist}
	done
}

function count_of_nodes()
{
	echo ${#NODE_ARRAY[@]}
}

function cpu_count_of_node()
{
	local node=${1:-0}
	local cpu_array=(${CPULIST_ARRAY[${node}]})
	echo ${#cpu_array[@]}
}

function count_of_cpus()
{
	local cpu_cnt=0
	local n
	for n in "${NODE_ARRAY[@]}"; do
		((cpu_cnt+=$(cpu_count_of_node ${n})))
	done
	echo ${cpu_cnt}
}

# input:  [NUMA node(0)] [max cpu count(all)] [from cpu index(0)] [next cpu index step(1)]
# output: <NUMA node> <cpu count> [cpu list from the node]
function cpulist_of_node()
{
	local node=${1:-0}
	local cpu_array=(${CPULIST_ARRAY[${node}]})
	local max_cpu_cnt=${2:-${#cpu_array[@]}}
	local cpu_idx=${3:-0}
	local cpu_idx_step=${4:-1}
	local cpulist=""
	local cpu_cnt=0
	local c
	for c in "${cpu_array[@]}"; do
		((cpu_idx < 0 || cpu_idx >= ${#cpu_array[@]})) && continue
		((cpu_cnt >= max_cpu_cnt)) && break
		((cpu_cnt++))
		cpulist=${cpulist:+${cpulist},}${cpu_array[${cpu_idx}]}
		((cpu_idx+=cpu_idx_step))
	done
	cpulist=$(compact_range_ex ${cpulist})
	echo "${node} ${cpu_cnt} ${cpulist}"
}

# input:  [NUMA node(0)] [cpu index range(null)]
# output: <NUMA node> <cpu count> [cpu list from the node]
function cpulist_of_node_ex()
{
	local node=${1:-0}
	local cpu_array=(${CPULIST_ARRAY[${node}]})
	local cpu_idx_range=($(expand_range ${2}))
	local cpulist=""
	local cpu_cnt=0
	local cpu_idx
	if ((${#cpu_array[@]} > 0)); then
		for cpu_idx in "${cpu_idx_range[@]}"; do
			((cpu_idx < 0 || cpu_idx >= ${#cpu_array[@]})) && continue
			((cpu_cnt++))
			cpulist=${cpulist:+${cpulist},}${cpu_array[${cpu_idx}]}
		done
		cpulist=$(compact_range_ex ${cpulist})
	fi
	echo "${node} ${cpu_cnt} ${cpulist}"
}

# input:  [max node count(all)] [from node index(0)] [next node index step(1)] [max cpu count per node(all)] [from cpu index per node(0)] [next cpu index step(1)]
# output: <node count> [node list] <cpu count> [cpu list]
function cpulist_of_nodes()
{
	local max_node_cnt=${1:-${#NODE_ARRAY[@]}}
	local node_idx=${2:-0}
	local node_idx_step=${3:-1}
	local cpu_cnt_per_node=${4}
	local cpu_idx_per_node=${5}
	local cpu_idx_step=${6}
	local node_cnt=0
	local nodelist
	local cpulist
	local cpu_cnt=0
	local n
	for n in "${NODE_ARRAY[@]}"; do
		((node_idx < 0 || node_idx >= ${#NODE_ARRAY[@]})) && continue
		((node_cnt >= max_node_cnt)) && break
		((node_cnt++))
		local node=${NODE_ARRAY[${node_idx}]}
		local tmp=($(cpulist_of_node ${node} ${cpu_cnt_per_node}) ${cpu_idx_per_node} ${cpu_idx_step})
		if ((${tmp[1]} > 0)); then
			nodelist=${nodelist:+${nodelist},}${tmp[0]}
			((cpu_cnt+=${tmp[1]}))
			cpulist=${cpulist:+${cpulist},}${tmp[2]}
		fi
		((node_idx+=node_idx_step))
	done
	nodelist=$(compact_range_ex ${nodelist})
	cpulist=$(compact_range_ex ${cpulist})
	echo "${node_cnt} ${nodelist} ${cpu_cnt} ${cpulist}"
}

# input:  [node index range(null)] [max cpu count per node(all)] [from cpu index per node(0)] [next cpu index step(1)]
# output: <node count> [node list] <cpu count> [cpu list]
function cpulist_of_nodes_ex()
{
	local node_idx_range=($(expand_range ${1}))
	local cpu_cnt_per_node=${2}
	local cpu_idx_per_node=${3}
	local cpu_idx_step=${4}
	local node_cnt=0
	local nodelist
	local cpulist
	local cpu_cnt=0
	local node_idx
	for node_idx in "${node_idx_range[@]}"; do
		((node_idx < 0 || node_idx >= ${#NODE_ARRAY[@]})) && continue
		((node_cnt++))
		local node=${NODE_ARRAY[${node_idx}]}
		local tmp=($(cpulist_of_node ${node} ${cpu_cnt_per_node}) ${cpu_idx_per_node} ${cpu_idx_step})
		if ((${tmp[1]} > 0)); then
			nodelist=${nodelist:+${nodelist},}${tmp[0]}
			((cpu_cnt+=${tmp[1]}))
			cpulist=${cpulist:+${cpulist},}${tmp[2]}
		fi
	done
	nodelist=$(compact_range_ex ${nodelist})
	cpulist=$(compact_range_ex ${cpulist})
	echo "${node_cnt} ${nodelist} ${cpu_cnt} ${cpulist}"
}

# input:  [node index range(null)] [cpu index range(null)]
# output: <node count> [node list] <cpu count> [cpu list]
function cpulist_of_nodes_range()
{
	local node_idx_range=($(expand_range ${1}))
	local cpu_idx_range="${2}"
	local node_cnt=0
	local nodelist
	local cpulist
	local cpu_cnt=0
	local node_idx
	for node_idx in "${node_idx_range[@]}"; do
		((node_idx < 0 || node_idx >= ${#NODE_ARRAY[@]})) && continue
		((node_cnt++))
		local node=${NODE_ARRAY[${node_idx}]}
		local tmp=($(cpulist_of_node_ex ${node} ${cpu_idx_range}))
		if ((${tmp[1]} > 0)); then
			nodelist=${nodelist:+${nodelist},}${tmp[0]}
			((cpu_cnt+=${tmp[1]}))
			cpulist=${cpulist:+${cpulist},}${tmp[2]}
		fi
	done
	nodelist=$(compact_range_ex ${nodelist})
	cpulist=$(compact_range_ex ${cpulist})
	echo "${node_cnt} ${nodelist} ${cpu_cnt} ${cpulist}"
}

# input: <dev class> <dev name>
# output: [NUMA node]
function node_of_dev()
{
	[[ -z "$1" || -z "$2" ]] && echo "" && return
	dev_class="/sys/class/$1/$2"
	[[ ! -e "${dev_class}" ]] && echo "" && return
	numa_node=$(readlink -f "${dev_class}" | grep -oE "/sys/devices/pci[0-9a-fA-F:/]+\.[0-9]+/")numa_node
	[[ ! -f "${numa_node}" ]] && echo "" && return
	cat "${numa_node}"
}

# input:  <net dev name>
# output: [NUMA node]
function node_of_net_dev()
{
	node_of_dev "net" "$1"
}

# input:  <block dev name>
# output: [NUMA node]
function node_of_block_dev()
{
	node_of_dev "block" "$1"
}

# =========== example begin ==========

function test_example_on_node()
{
	echo "node: ${1}, cpu count: ${2}, cpu list: ${3}"
}

function test_example_on_nodes()
{
	if (($# == 4)); then
		echo "node count: ${1}, node list: ${2}, cpu count: ${3}, cpu list: ${4}"
	else
		echo "node count: ${1}, node list: , cpu count: ${2}, cpu list: "
	fi
}

function test_example_0()
{
	echo "${FUNCNAME[0]}:"
	echo "count of nodes: $(count_of_nodes)"
	echo "cpu count of/per node(0): $(cpu_count_of_node)"
	echo "count of cpus: $(count_of_cpus)"
	echo "NODE_ARRAY: ${NODE_ARRAY[@]}"
	echo "CPULIST_ARRAY: ${CPULIST_ARRAY[@]}"
	local n
	for n in "${NODE_ARRAY[@]}"; do
		echo "NODE ${n} CPULIST: ${CPULIST_ARRAY[${n}]}"
	done
}

function test_example_1()
{
	echo "${FUNCNAME[0]}:"
	local n
	local k
	for n in "${NODE_ARRAY[@]}"; do
		for ((k = 1; k <= $(cpu_count_of_node ${n}); k++)); do
			local cpulist=$(cpulist_of_node ${n} ${k} 0)
			test_example_on_node ${cpulist}
		done
	done
}

function test_example_2()
{
	echo "${FUNCNAME[0]}:"
	local i
	local k
	for ((i = 0; i < $(count_of_nodes); i++)); do
		for ((k = 1; k <= $(cpu_count_of_node ${NODE_ARRAY[${i}]}); k++)); do
			local cpulist=$(cpulist_of_nodes $((i + 1)) 0 1 ${k} 0)
			test_example_on_nodes ${cpulist}
		done
	done
}

function test_example_3()
{
	echo "${FUNCNAME[0]}:"
	local k
	for ((k = 1; k <= $(cpu_count_of_node); k++)); do
		local cpulist=$(cpulist_of_nodes $(count_of_nodes) 0 1 ${k} 0)
		test_example_on_nodes ${cpulist}
	done
}

function test_example_4()
{
	echo "${FUNCNAME[0]}:"
	test_example_on_node $(cpulist_of_node 1 2 3)
}

function test_example_5()
{
	echo "${FUNCNAME[0]}:"
	test_example_on_nodes $(cpulist_of_nodes 4 2 2 4 2)
}

function test_example_6()
{
	echo "${FUNCNAME[0]}:"
	test_example_on_node $(cpulist_of_node_ex 2 0,2,4,6-$(($(cpu_count_of_node 1) - 1)))
}

function test_example_7()
{
	echo "${FUNCNAME[0]}:"
	test_example_on_nodes $(cpulist_of_nodes_ex 0-$(($(count_of_nodes) - 1)) 4 0)
}

function test_example_8()
{
	echo "${FUNCNAME[0]}:"
	test_example_on_nodes $(cpulist_of_nodes_range 0-$(($(count_of_nodes) - 1)) 0,2,4,1-$(($(cpu_count_of_node 0) - 1)))
}

function test_example_9()
{
	echo "${FUNCNAME[0]}:"
	local n=0
	local cpu_array=(${CPULIST_ARRAY[${n}]})
	local c
	for c in "${cpu_array[@]}"; do
		test_example_on_node $(cpulist_of_node ${n} ${#cpu_array[@]} ${c})
	done
}

function test_example_10()
{
	echo "${FUNCNAME[0]}:"
	test_example_on_node $(cpulist_of_node)
	test_example_on_node $(cpulist_of_node_ex)
	test_example_on_nodes $(cpulist_of_nodes)
	test_example_on_nodes $(cpulist_of_nodes_ex)
	test_example_on_nodes $(cpulist_of_nodes_range)
}

function test_example_11()
{
	echo "${FUNCNAME[0]}:"
	local dev
	for dev in `lsblk -l | grep -oE "^[ns][^ ]*"`; do
		echo "NUMA node of $dev: $(node_of_block_dev ${dev})"
	done
}

function test_example_12()
{
	echo "${FUNCNAME[0]}:"
	local dev
	for dev in `ip link show | grep -E "^[0-9]+:" | awk -F': ' '{print $2}'`; do
		echo "NUMA node of $dev: $(node_of_net_dev ${dev})"
	done
}

function test_example_13()
{
	echo "${FUNCNAME[0]}:"
	index_range
	index_range 10 0 1
	index_range 10 5 2
	index_range 10 5 -1
}

function test_example_14()
{
	echo "${FUNCNAME[0]}:"
	index_range_ex
	index_range_ex 10 0 1
	index_range_ex 10 5 2
	index_range_ex 10 5 -1
	index_range_ex 10 10 2
	index_range_ex 0 79 1
	index_range_ex 0 80 1
	index_range_ex 0 80 2
	index_range_ex 0 159 1
	index_range_ex 0 160 1
	index_range_ex 0 160 2
	index_range_ex 80 80 2
	index_range_ex 80 80 2
	local idx=($(index_range_ex 0 160 1))
	echo ${idx[0]}
	echo ${idx[1]}
	echo ${idx[1]//,/&}
	echo `compact_range ${idx[1]}`
	local idx=($(index_range_ex 0 160 2))
	echo ${idx[0]}
	echo ${idx[1]}
	echo `compact_range ${idx[1]}`
}

function run_test_examples()
{
	init_cpu_nodes
	local exams=($(expand_range $@))
	local examples=($(declare -F | grep -oE "test_example_[0-9]+$" | sort -V))
	printf "%s\n%s\n" "All test examples(${#examples[@]}):" "${examples[@]}"
	local t
	if ((${#exams[@]} > 0)); then
		local n
		for n in "${exams[@]}"; do
			t=${examples[((${n} % ${#examples[@]}))]}
			echo
			${t}
		done
	else
		for t in "${examples[@]}"; do
			echo
			${t}
		done
	fi
}

if [[ "$1" =~ example[s]{0,1}$ ]]; then
	shift
	run_test_examples $@
fi

# =========== example end ==========
